#!/usr/bin/env python
# -*- coding: UTF-8 -*-

import os, re, pickle
from struct import *
from time import time, strftime, gmtime

def help():
	print '''
Search
>>> from cnlstk.search import Search
>>>
>>> #設定搜尋的文本檔案、索引檔、及LB檔案資訊
>>> Fobj = Search('Text_f', 'sch', 'Text_f_SA', 'Text_lb')
>>> line = u'搜尋字串'
>>> 
>>> #設定搜尋的檔案範圍 (預設為搜尋全部檔案)
>>> Fobj.setFilter(('h001','h006','h056','h080'))
>>> Fobj.getFilter()
>>>
>>> #清除搜尋的檔案範圍 (預設為搜尋全部檔案)
>>> Fobj.resetFilter()
>>> Fobj.getFilter()
>>>
>>> # getConcordance
>>> rst = Fobj.getConcordance(line)
>>> for k in rst:
>>> 	print k[0], k[1]
>>> 
>>> # getNextWords
>>> rst = Fobj.getNextWords(line, 1)
>>> for k in rst:
>>> 	print k[0], k[1]
>>> 
>>> # getNextWordCounts
>>> rst = Fobj.getNextWordCounts(line)
>>> print rst
>>> 
'''

class Search:
	'''建立與檢索 Suffix Array 的 class'''
	
	def __init__( self, fname, opt, suffix, lbindex ):
		'''
		[ fname ] = The main text file. 
		[ suffix ] = Suffix Array of fname
		[ lbindex ] = lb index of fname
		[ opt ] = 'sch' or 'schbk'
		'''
		self.F = open( fname, 'rb' )			##開啟主文字檔##
		self.F_len = os.stat( fname )[6]		##主文字檔大小
		self.File_main = fname					##主要檔的檔名##
		
		self.S = open( suffix, 'rb' )			##主要檔案的 Suffix Array##
		self.S_len = os.stat( suffix )[6]		##Suffix Array 的檔案大小##
		self.Suffix_A = suffix					##主要 Suffix Array 的檔名##
		if self.S_len > 1000000:				#1mb 大小以上的 Suffix Array 才建索引 cache
			self.cache_p = [0]*1024				##紀錄 Suffix Array 位置的 cache##
			self.cache_w = [""]*1024			##紀錄 self.cache_p 對應回 self.S 再對應回 self.F 的內容字串##
			self.__getCache( 0, self.S_len/4, 1 )

		self.opt = opt							#正向檔或反向檔
		
		pkl_f = open( lbindex, 'rb' )
		self.section = pickle.load( pkl_f )		#範圍資訊索引 Dic
		self.seckey = self.section.keys()
		self.seckey.sort()
		pkl_f.close()
		
		self.filter = []						#縮小範圍依據
		
	def __getCache( self, i, j, c ):
		"""Create 2 cache lists for search.
			[ i ] = the start position
			[ j ] = the end position
			[ c ] = position of cache list"""
			
		if c < 1024:
			k = ( i+j )/2		##中間位置 ( k ) = 首尾位置 ( i, j ) 除以 2##
			self.cache_p[c] = k
			self.S.seek( k*4 )
			ofst = unpack( 'I', self.S.read(4) )[0]	##提出 self.S 實際值##
			w = self.__getWords( ofst, 32 )		##取出 self.F 32個字##
			self.cache_w[c] = w
			
			self.__getCache( i, k, c*2 )
			self.__getCache( k, j, c*2+1)

	def __getWords( self, ofst, length ):
		'''
		輸入 offset + 要取的字串長度
		回傳 offset 開始向後（右）要取的長度字串
		'''
		line = ''
		self.F.seek( ofst )
		bytes = (length + length/2) * 2		#轉成 bytes 數，比需要長度多取一半
		try:
			line = self.F.read(bytes).decode('utf16')
		except:
			self.F.seek( ofst )
			line = self.F.read(bytes+2).decode('utf16')
		return line[:length]
		
	def __getsdroW( self, ofst, length ):
		'''
		輸入 offset + 要取的字串長度
		回傳 offset 向前（左）要取的長度字串
		'''
		line = ''
		for k in range(length):
			if ofst - 2 >= 2:			#還沒到檔頭
				try:
					self.F.seek( ofst-2 )
					line = self.F.read(2).decode('utf16') + line
					ofst = ofst - 2
				except:
					self.F.seek( ofst-4 )
					line = self.F.read(4).decode('utf16') + line
					ofst = ofst - 4
			else:
				break
#				line = u'　' + line
		return line

####### Search #######
	def __backwardStr( self, strg ):
		"""Return a backward string of [ strg ]"""
		line = ""		##儲存字串反過來的值##
		for k in strg:
			line = k + line
		return line
	
	def getLB(self, offset):
		'''
		輸入 offset, 回傳該 offset 所在檔案
		'''
		pre = ''
		if self.opt == 'sch':
			for k in self.seckey:
				if offset < self.section[k][0]:
					break
				pre = k
		elif self.opt == 'schbk':
			for k in self.seckey:
				if offset > self.section[k][1]:
					break
				pre = k
		return pre

	def getRange(self, lb):
		'''
		輸入行號，回傳該行首尾 offset
		'''
		ans_s = 0
		ans_e = 0
		i = 0
		j = self.cache_line_len/28			#長度相同所以沒換
		while i <= j:
			p = (i+j)/2
			self.cache_line.seek(p*28)
			position = unpack( "17sII", self.cache_line.read(28) )
			#position = ('T01n0001_p0001a01', 24, 290876352)
			if self.opt == 'schbk':
				post = position[2]
			else:
				post = position[1]
			
			if lb < position[0]:
				j = p-1
				ans_e = post
#				print '<', ans_s, ans_e
			else:
				i = p+1
				ans_s = post
#				print '>=', ans_s, ans_e
		return ans_s, ans_e
		
	def getRangeWords( self, start, end ):
		if (end <= start) or (start < 2) or (end < 2):
			rst = 'wrong range!'
		elif self.opt == 'schbk':
			self.F.seek(start)
			try:
				rst = self.F.read( end+2 - start ).decode('utf16')[1:]
			except:
				rst = self.F.read( end+4 - start ).decode('utf16')[1:]
		else:
			self.F.seek( start )
			rst = self.F.read( end - start ).decode('utf16')
		return re.sub(r'^\n+|\n+$', '', rst)
		
	def getOfsetList(self, SAspot, times):
		'''
		回傳字串 offset list
		字串結果照 offset 順序排列的陣列，為查字串在哪些經文中用。
		'''
		'''
		D = {}
		self.S.seek(SAspot)
		for k in range(times):
			D[k] = unpack( 'I', self.S.read(4) )[0]
		'''
		D = []
		self.S.seek(SAspot)
		for k in range(times):
			D.append(unpack( 'I', self.S.read(4) )[0])
		D.sort()
		return D
			
	def getText2(self, offset, bw, aw):
		'''
		給定 offset 取回全文與行號資訊
		'''
		lb = self.getLB(offset)
		tmp = self.__getsdroW( offset, bw ) + self.__getWords( offset, aw )
		if self.opt == 'schbk':
			tmp = self.__backwardStr( tmp )
		return (tmp, lb)
		
#### Set search range ####	
	def getSecList(self):
		'''
		取得限定範圍用的 keys
		'''
		return self.seckey
		
	def setFilter(self, L):
		'''
		傳入限定的範圍陣列（以 lb keys 為依據）
		設定限定範圍的參數 self.filter
		'''
		for k in range(len(self.seckey)):
			if k+1 >= len(self.seckey):
				fe = self.F_len
				bs = 2
			else:
				fe = self.section[self.seckey[k+1]][0]
				bs = self.section[self.seckey[k+1]][1]
			fd = (self.section[self.seckey[k]][0], fe)
			bk = (bs, self.section[self.seckey[k]][1])
			if self.seckey[k] in L:
				self.filter.append((fd, bk))
		return 'filter set'
	
	def getFilter(self):
		return self.filter
		
	def resetFilter(self):
		self.filter = []

#### Next word count ####	
	def getNextWords(self, kw, kwlen, opt=0):
		'''
		kw: search string
		kwlen: how many next words
		opt: 0 only return (next words)
		opt: 1 return (next words, frequency)
		'''
		n = self.run(kw)
		times = n[0]
		SAspot = n[1] 
		
		nL = []		#有索引的 Next Word 的記錄
		ct = 0
#		bp = 'no'	#記錄有沒有斷點符號
		while times > ct:
			self.S.seek( SAspot )
			offset = unpack( 'I', self.S.read(4) )[0]
#			print offset, len(kw), kwlen

			n = self.__getWords( offset, len(kw)+kwlen )	# kw='A', kwlen='2', n='ABC', m='BC'
			m = n[-kwlen:]

			#這一段好像是偵測錯誤用
			if self.opt == 'schbk':
				n = self.__backwardStr(n)
			rst = self.run(n)
			if rst[0] == 0:
				nL = [('none', 0)]
				break
			
			if opt == 0:
#				if (m[0] >= u'\u3040' and m[0] <= u'\uFAFF') or m[0] > u'\uFFEF' or m[0] == u'\u2E84':
				nL.append( (m, rst[0]) )
			else:
				nL.append(m)

			ct += rst[0]
			SAspot = SAspot + rst[0]*4
#			print m, rst[0], ct, '/', times
#			raw_input()
		return nL

	def getNextWordCounts(self, kw, kwlen=1):
		'''
		與 getNextWords() 輸出不同，抽詞用
		return [幾種（不含斷點）, 斷點數, 種類中最多的數量]
		'''
		n = self.run(kw)
		times = n[0]
		SAspot = n[1]
		
		nL = []		#有索引的 Next Word 的記錄
		clt = 0		#換行或其他非文字符號（沒見索引的字碼）的次數
		ct = 0
		ii = 0		#右側種類數
		max = 0		#ii中最多者
		bb = 0		#斷點數
		while times > ct:
			self.S.seek( SAspot )
			offset = unpack( 'I', self.S.read(4) )[0]
			n = self.__getWords( offset, len(kw)+kwlen )
#			print 'search', n.replace('\n', '-')
			m = n[-kwlen:]

			if self.opt == 'schbk':
				n = self.__backwardStr(n)
			rst = self.run(n)
			if rst[0] == 0:
				ii = bb = max = -1
				break
				
			if (m[0] >= u'\u3040' and m[0] <= u'\uFAFF') or m[0] > u'\uFFEF' or m[0] == u'\u2E84':
				ii += 1
				if rst[0] > max:
					max = rst[0]
			else:
				bb += rst[0]
#			nL.append( (m, rst[0]) )
			ct += rst[0]
			SAspot = SAspot + rst[0]*4
#			print m, rst[0], ct, '/', times
#			raw_input()
		return {'distinct next word count':ii, 'max no. of next words':max, 'char count of non-Chinese':bb}

#### Concordance ####
	def getText(self, offset, bw, aw):
		'''
		傳入 offset 及字串長度
		回傳所需字串長度 + 出現檔案名稱
		'''
		lb = self.getLB(offset)
		tmp = self.__getsdroW( offset, bw ) + self.__getWords( offset, aw )
		if self.opt == 'schbk':
			tmp = self.__backwardStr( tmp )							
		line = tmp.replace( "\n", u"\u204B " )
		return (line, lb)
	
	def getConcordance(self, KW, addw=10):
		'''
		輸入 索引名、字串、開始筆、筆數（一次最高 30 筆）、結果的前後長度（內定前後各 10 個字）
		回傳字串與 CBETA 行號
		'''
		n = self.run(KW)
		if n[0] == 0:
			return ('No records', '')

		L = []
		self.S.seek( n[1] )
		if self.filter == []:
			for k in range(n[0]):
				offset = unpack( 'I', self.S.read(4) )[0]
				L.append( self.getText(offset, addw, len(KW)+addw) )
		else:
			if self.opt == 'sch':
				ck = 0		#self.filter 所有位置 0 的為正向 ofset 範圍
			else:
				ck = 1		#反向為 1

			for k in range(n[0]):
				offset = unpack( 'I', self.S.read(4) )[0]
				for k in self.filter:
					if offset >= k[ck][0] and offset <= k[ck][1]:
						L.append( self.getText(offset, addw, len(KW)+addw) )
		return L
		
#### Binary search ####
	def run( self, KW ):
		'''
		輸入欲查詢的字串
		回傳 (該字串出現次數, SuffixArray 檔的起始 offset)
		無此字串時回傳 (0, -1)
		'''
		
		if self.opt == 'schbk':
			KW = self.__backwardStr( KW )
			
		kw_len = len(KW)			##欲搜尋字串的字數##
		S_end = self.S_len/4		##suffix array 的長度(個數)##
		i = 0							##binary search 的起始位置##
		j = self.S_len/4				##binary search 的結束位置##
#		print "search for %s, len: %d, sa: %d, sa_s: %d, sa_e: %d" % (KW, kw_len, self.S_len, i, j)
		if self.S_len > 1000000:	#有 cache
			ck = 1											##cache lists 位置的參數##
			while ck < 1024:
				k = self.cache_p[ck]					##位置 cache 的參數##
				w = self.cache_w[ck][0:len(KW)]	##字串 cache 的參數##
				
				"""#如果 cache 字串長度小於欲搜尋的字串長度. suffix array > 1mb 才做 cache, 不會遇到這種狀況
				try:
					w = self.cache_w[ck][0:len(KW)]
				except:
					w = self.cache_w[ck]"""
				
				if KW < w:
					j = k-1
					ck = ck*2
				elif KW > w:
					i = k+1
					ck = ck*2+1
				else:
					break	#比對到直接跳出, 剛剛的 i, j 傳給下一個迴圈
#			print i, j
#			raw_input()

		ans = ( 0, -1 )		##搜尋結果參數##
		while i <= j:
			k = ( i+j )/2
			if k == S_end:	#最後一個位置, 後面沒有資料
				break
			self.S.seek( k*4 )
			offset = unpack( "I", self.S.read(4) )[0]	##suffix array 中該位置紀錄的 offset (for self.F)##
			w = self.__getWords( offset, kw_len )		##每個offset向後抓KW同長度的字串##
#			print w
#			raw_input()	
			if KW == w:	#比對到第一個
				sa_start_p = k		##紀錄出現該字串 suffix array 的開始與結束位置##
				sa_end_p = k	
				
				k_bk = k
				while i <= k_bk:	#前半
					ks = ( i+k_bk )/2
					self.S.seek( ks*4 )
					offset = unpack( "I", self.S.read(4) )[0]
					w = self.__getWords( offset, kw_len )	#到每個offset向後抓KW同長度的字串來比
					if KW == w:
						sa_start_p = ks
						k_bk = ks-1
					elif KW != w:
						i = ks+1
						sa_start_p = i	#目前的 ks 確定不是, 下一個才可能是
				while k <= j:	#後半
					ks = ( k+j )/2
					if ks == S_end:
						break
					self.S.seek( ks*4 )
					offset = unpack( "I", self.S.read(4) )[0]
					w = self.__getWords( offset, kw_len )	#到每個offset向後抓KW同長度的字串來比
					if KW == w:
						sa_end_p = ks
						k = ks+1
					elif KW != w:
						j = ks-1
						sa_end_p = j
				c = sa_end_p - sa_start_p + 1		##搜尋到[ KW ]的總數 = 開始位置 - 結束位置 + 1 ##
				ans = ( c, sa_start_p*4 )			##搜尋結果. 回傳次數與起始的 suffix array 檔的 bytes數##
				break
				
			elif KW < w:
				j = k-1
			else:
				i = k+1
		return ans
	
	
	